#pragma once

#include <algorithm>
#include <shared_mutex>
#include <unordered_set>

#include <butil/string_printf.h>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <json2pb/json_to_pb.h>
#include <json2pb/pb_to_json.h>

#include "hawking/indexlib/indexlib_service.pb.h"

#include "index/index.h"
#include "index/attribute_index/attribute_indexer.h"
#include "index/inverted_index/inverted_indexer.h"
#include "index/primary_key_index/primary_key_indexer.h"
#include "index/deletion_index/deletion_index_indexer.h"
#include "util/hashmap/hashset.h"
#include "util/status.h"

namespace hawking {
namespace indexlib {

DECLARE_uint32(max_fetch_doc_count);

inline bool IsSameSchema(
    std::shared_ptr<IndexlibIndexConfig> a, std::shared_ptr<IndexlibIndexConfig> b) {
    return (a && b) ?
        (google::protobuf::util::MessageDifferencer::Equals(*a, *b))
        : ((!a && !b) ? true : false);
}

class TombStone {
public:
    explicit TombStone(int32_t init_size) {
        running_.store(true, std::memory_order_relaxed);
        tomb_stone_ = std::make_unique<HashSet<uint64_t>>(init_size);
        tomb_doc_ids_ = std::make_unique<HashSet<DocId>>(init_size);
    }

    void InsertTomb(uint64_t pk, DocId doc_id) {
        bool running = running_.load(std::memory_order_acquire);
        if (running) {
            tomb_stone_.emplace(pk);
            tomb_doc_ids_->Insert(doc_id);
        }
    }

    void Stop() {
        running_.store(false, std::memory_order_acq_rel);
    }

    bool ExistDocId(DocId doc_id) const {
        return tomb_doc_ids_ && tomb_doc_ids_->Exist(doc_id);
    }

    const std::unique_ptr<HashSet<DocId>>& Pks() const {
        return tomb_stone_;
    }

private:
    std::atomic<bool> running_;
    std::unique_ptr<HashSet<uint64_t>> tomb_stone_;
    std::unique_ptr<HashSet<DocId>> tomb_doc_ids_;
};

class Segment {
public:
    enum class SegmentState {
        Idle,
        Writing,
        DumpingOrMerging
    };

    Segment() {
        is_dirty_.store(false, std::memory_order_relaxed);
        ResetSegmentState();
    }
    virtual ~Segment() = default;

    Status Query(
        const std::string& invert_field_nane,
        uint64_t term_key,
        uint32_t topk,
        const google::protobuf::RepeatedPtrField<std::string>& select_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& filter_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& where_fields,
        IndexlibRetrieveResponse* resp);
    
    Status QueryByPk(
        const google::protobuf::RepeatedField<int64_t>& pks,
        const google::protobuf::RepeatedPtrField<std::string>& select_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& filter_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& where_fields,
        IndexlibRetrieveResponse* resp);

    ReturnValue<bool> FilterCheck(
        int64_t pk,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& filter_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& where_fields);
    ReturnValue<bool> FilterCheck(
        int64_t pk, const util::FilterColumn& field, bool is_where);
    
    ReturnValue<bool> FillLocaldata(
        int64_t pk,
        google::protobuf::RepeatedPtrField<util::SelectColumn>* select_fields);
    ReturnValue<bool> FillLocaldata(
        int64_t pk,
        util::SelectColumn* select_field);
    ReturnValue<bool> FillLocaldata(
        int64_t pk, int32_t field_id, util::FieldValue* field_value);

    Status Delete(uint64_t pk) override {
        ReturnValue<DocId> ret = pk_indexer_->Lookup(pk);
        if (ret.Code() == Status::OK) {
            if (CouldWrite()) {
                deletion_indexer_->Delete(ret.Value());
                ResetSegmentState();
            } else {
                InsertTomb_(pk, ret.Value());
            }
            return Status::OK;
        }

        return Status::NotFound;
    }

    ReturnValue<DocId> Lookup(uint64_t pk) const override {
        return pk_indexer_ ? pk_indexer_->Lookup(pk) : {Status::NotFound, 0};
    }

    virtual bool ShouldDump() const = 0;

    uint32_t DocCount() const override {
        return pk_indexer_->DocCount() - deletion_indexer_->GetDeletedDocCount();
    }

    // attribute fields update
    util::HawkingRecallRetCode UpdateDocument(
        const IndexlibMessage& message, uint64_t pk,
        std::shared_ptr<NormalDocument> doc, DocId doc_id) {
        if (unlikely(!need_increase_)) {
            return util::NOT_SUPPORT;
        }
        
        if (!CouldWrite()) {
            return util::SEGMENT_DUMPING_OR_MERGING;
        }

        util::HawkingRecallRetCode ret = UpdateDocument_(message, pk, doc, doc_id);
        ResetSegmentState();
        return ret;
    }

    virtual bool Dump(const std::string& segment_dir) = 0;

    std::unique_ptr<HashSet<DocId>>& GetTombStone() {
        return tomb_stone_;
    }

    std::shared_ptr<IndexlibIndexConfig> Schema() {
        return schema_;
    }

    std::shared_ptr<DocumentCreator> DocCreator() {
        return document_creator_;
    }

    int64_t TTL() const {
        return schema_.life_cycle();
    }

    bool CouldWrite() {
        SegmentState segment_state = SegmentState::Idle;
        return segment_state_flag_.compare_exchange_strong(
            segment_state, SegmentState::Writing);
    }

    void WaitForWriteFinish() {
        while (1) {
            SegmentState segment_state = SegmentState::Idle;
            if (segment_state_flag_.compare_exchange_strong(
                    segment_state, SegmentState::DumpingOrMerging)) {
                return;
            } else {
                bthread_yield();
            }
        }
    }

    void ResetSegmentState() {
        segment_state_flag_.store(SegmentState::Idle, std::memory_order_release);
    }

    bool IsDirty() const {
        return is_dirty_.load(std::memory_order_relaxed);
    }

    bool GetFieldIdAndTypedByName(
        const std::string& field_name,
        int32_t* field_id,
        util::ColumnType* field_type) const {
        return document_creator_->FieldName2IdAndType(field_name, field_id, field_type);
    }

    bool HasInvertedIndexer() {
        return !inverted_indexers_.empty();
    }

    std::shared_ptr<InvertedIndexer> GetInvertedIndexer(int32_t field_id) {
        auto iter = inverted_indexers_.find(std::make_pair(field_name, field_type));
        if (iter == inverted_indexers_.end()) {
            return nullptr;
        }

        return iter->second;
    }
    std::shared_ptr<AttributeIndexer> GetAttributeIndexer(int32_t field_id) {
        auto iter = attribute_indexers_.find(field_id);
        if (iter == attribute_indexers_.end()) {
            return nullptr;
        }

        return iter->second;
    }
    std::shared_ptr<PrimaryKeyIndexer> GetPkIndexer() {
        return pk_indexer_;
    }
    std::shared_ptr<DeletionMapIndexer> GetDeletionIndexer() {
        return deletion_indexer_;
    }

protected:
    util::HawkingRecallRetCode UpdateInvertedIndex_();
    util::HawkingRecallRetCode UpdateAttributeIndex_();

    util::HawkingRecallRetCode UpdateDocument_(
        const IndexlibMessage& message, uint64_t pk,
        std::shared_ptr<NormalDocument> doc, DocId doc_id);

    void TombStoneFilter_(SegmentContext* segment_ctx) const;

    void InsertTomb_(uint64_t pk, DocId doc_id) {
        if (tomb_stone_) {
            tomb_stone_->InsertTomb(pk, doc_id);
        }
    }

    void SetDirty_() {
        is_dirty_.store(false, std::memory_order_release);
    }

    std::shared_ptr<NormalDocument> CreateNormalDocument(
        const IndexlibMessage& message) const {
        return document_creator_->CreateNormalDocument(message);
    }

    void ObsoletedFilter_(
        std::unordered_map<DocId, DocId>* old2new_docid_map) const;

    // building/dumping is from MemSegment::Init, disk is from load from metafile
    std::shared_ptr<IndexlibIndexConfig> schema_;
    std::shared_ptr<DocumentCreator> document_creator_;

    // state machine
    std::atomic<SegmentState> segment_state_flag_;

    // record deleted doc by del/update during dumping/merging
    std::unique_ptr<TombStone> tomb_stone_;
    bool need_increase_ = false;

    // is modified
    std::atomic<bool> is_dirty_;

    // indexes
    std::unordered_map<
        int32_t field_id,
        std::shared_ptr<InvertedIndexer>> inverted_indexers_;
    std::unordered_map<
        int32_t field_id,
        std::shared_ptr<AttributeIndexer>> attribute_indexers_;
    std::shared_ptr<PrimaryKeyIndexer> pk_indexer_;
    std::shared_ptr<DeletionMapIndexer> deletion_indexer_;

private:
    void FilterDeletedDocs_(SegmentContext* segment_ctx) const;

    bool FilterDeletedDoc_(DocId doc_id) const;

    void MoveToResponse_(
        SegmentContext* segment_ctx,
        const google::protobuf::RepeatedPtrField<std::string>& select_fields,
        IndexlibRetrieveResponse* resp) const;

    Status ReadData_(
        std::shared_ptr<AttributeIndexer> attr_indexer,
        SegmentContext* segment_ctx,
        util::ColumnType field_type,
        int32_t field_id) override;

    void DoSelect_(
        const google::protobuf::RepeatedPtrField<std::string>& select_fields,
        SegmentContext* segment_ctx);
    
    void DoSelect_(
        DocId doc_id,
        google::protobuf::RepeatedPtrField<util::SelectColumn>* select_fields);
    
    void DoSelect_(
        DocId doc_id, int32_t field_id, util::FieldValue* field_value);

    void DoFilter_(
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& filter_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& where_fields,
        SegmentContext* segment_ctx);
    void DoFilter_(
        const util::FilterColumn& field, SegmentContext* segment_ctx, bool is_where);

    void FilterByTTL_(SegmentContext* segment_ctx);

    Status CheckValid_(
        const std::string& field_name,
        int32_t* field_id, util::ColumnType* field_type,
        std::shared_ptr<AttributeIndexer>& attr_indexer) const;
    Status CheckValid_(
        int32_t field_id, util::ColumnType* field_type,
        std::shared_ptr<AttributeIndexer>& attr_indexer) const;

    void FilterAndSelect_(
        SegmentContext* segment_ctx,
        const google::protobuf::RepeatedPtrField<std::string> select_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn> filter_fields,
        const google::protobuf::RepeatedPtrField<util::FilterColumn>& where_fields,
        IndexlibRetrieveResponse* resp);

    util::HawkingRecallRetCode ExcludeFakeModify_(
        const IndexlibMessage& message,
        DocId doc_id,
        SegmentContext* segment_ctx,
        std::unordered_set<int32_t>* exclude_field_ids) const;

    bool CheckIsSameFieldValue_(
        const SegmentContext& segment_ctx, int32_t field_id,
        util::ColumnType field_type, const FieldValue& field_value) const;

    bool CalcTerm_(
        int32_t field_id, const SegmentContext& segment_ctx, uint64_t* term) const;

    void CheckInvertedFieldModified_(
        const IndexlibMessage& message,
        std::vector<std::pair<int32_t, const util::FieldValue&>>* inverted_fields) const;
};

}
}